package com.algorithm.Tree;


public class RBTree<E> extends BinarySearchTree<E>{
    private static final boolean RED = true;
    private static final boolean BLACK = false;

    @Override
    protected void afterAdd(Node<E> node) {
        Node<E> parent = node.parent;
        if (parent == null) {
            black(node);
            return;
        }

        if (isBlack(parent)) return;

        Node<E> uncle = parent.sibling();
        Node<E> grand = red(parent.parent);
        if (isRed(uncle)) {//r-b-r 需要溢出
            black(parent);
            black(uncle);
            afterAdd(grand);
            return;
        }

        if (parent.isLeftChild()) {
            if (node.isLeftChild()) {
                black(parent);
            }else {
                black(node);
                rotatedLeft(parent);
            }
            rotatedRight(grand);
        }else {
            if (node.isLeftChild()) {
                black(node);
                rotatedRight(parent);
            }else {
                black(parent);
            }
            rotatedLeft(grand);
        }
    }

    /**
     * r-b-r r-b b-r b
     */
    @Override
    protected void afterRemove(Node<E> node) {
        if (isRed(node)) {
            black(node);
            return;
        }

        Node<E> parent = node.parent;
        if (parent == null) return;

        boolean left = parent.left == null || node.isLeftChild();
        Node<E> sibling = left ? parent.right : parent.left;
        if (left) {
            if (isRed(sibling)) {
                red(parent);
                black(sibling);
                rotatedLeft(parent);
                sibling = parent.right;
            }

            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                boolean bNeedRet = isBlack(parent);
                red(sibling);
                black(parent);
                if (bNeedRet) {
                    afterRemove(parent);
                }
            }else {
                if (isBlack(sibling.right)) {
                    rotatedRight(sibling);
                    sibling = parent.right;
                }
                color(sibling,colorOf(parent));
                black(parent);
                black(sibling.right);
                rotatedLeft(parent);
            }
        }else {
            if (isRed(sibling)) {
                red(parent);
                black(sibling);
                rotatedRight(parent);
                sibling = parent.left;
            }

            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                boolean bNeedRet = isBlack(parent);
                red(sibling);
                black(parent);
                if (bNeedRet) {
                    afterRemove(parent);
                }
            }else {
                if (isBlack(sibling.left)) {
                    rotatedLeft(sibling);
                    sibling = parent.left;
                }
                color(sibling,colorOf(parent));
                black(parent);
                black(sibling.left);
                rotatedRight(parent);
            }
        }
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new RBNode<>(element,parent);
    }

    private boolean colorOf(Node<E> node) {
        return node == null ? BLACK : ((RBNode) node).color;
    }

    private boolean isRed(Node<E> node) {
        return colorOf(node) == RED;
    }

    private boolean isBlack(Node<E> node) {
        return colorOf(node) == BLACK;
    }

    private Node<E> color(Node<E> node,boolean color) {
        if (node == null) return null;
        ((RBNode) node).color = color;
        return node;
    }

    private Node<E> red(Node<E> node) {
        return color(node, RED);
    }

    private Node<E> black(Node<E> node) {
        return color(node,BLACK);
    }

    private static class RBNode<E> extends Node<E> {
        boolean color = RED;

        RBNode(E element, Node<E> parent) {
            super(element, parent);
        }

        @Override
        public String toString() {
            return element + (color == RED ? "_RED" : "");
        }
    }
}
